#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Software: PyCharm
# @Version : Python-
# @Author  : Shengji He
# @Email   : hsjbit@163.com
# @File    : SumDistanceTree.py
# @Time    : 2020/10/6 14:09
# @Description:
from typing import List


class Solution:
    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        """
        An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges are given.
        
        The ith edge connects nodes edges[i][0] and edges[i][1] together.
        
        Return a list ans, where ans[i] is the sum of the distances between node i and all other nodes.
        
        Example 1:
        
            Input: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
            Output: [8,12,6,10,10,10]
            Explanation:
                Here is a diagram of the given tree:
                  0
                 / \
                1   2
                   /|\
                  3 4 5
            We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
            equals 1 + 1 + 2 + 2 + 2 = 8.  Hence, answer[0] = 8, and so on.

        Note: 1 <= N <= 10000

        :param N:
        :param edges:
        :return:
        """
        tree = [[] for _ in range(N)]
        for baba, erza in edges:
            tree[baba].append(erza)
            tree[erza].append(baba)
        depth = [0 for _ in range(N)]
        count = [0 for _ in range(N)]

        def dfsForDepthAndCount(baba, father):
            count[baba] = 1
            for erza in tree[baba]:
                if erza != father:
                    depth[erza] = depth[baba] + 1
                    dfsForDepthAndCount(erza, baba)
                    count[baba] += count[erza]

        dfsForDepthAndCount(0, -1)
        answer = [0 for _ in range(N)]
        answer[0] = sum(depth)

        def dfs_for_answer(baba, father):
            for erza in tree[baba]:
                if erza != father:
                    answer[erza] = answer[baba] + N - 2 * count[erza]
                    dfs_for_answer(erza, baba)

        dfs_for_answer(0, -1)
        return answer


if __name__ == '__main__':
    S = Solution()
    N = 6
    edges = [[0, 1], [0, 2], [2, 3], [2, 4], [2, 5]]
    print(S.sumOfDistancesInTree(N, edges))
    print('done')
